import pandas as pd
from pyecharts.charts import HeatMap
from pyecharts import options as opts

# 数据文件地址
csv_file = r"data\数据.csv"  # 替换为你的文件路径

# 读取 CSV 文件，使用制表符分隔
df = pd.read_csv(csv_file, sep="\t")

# 确保列名无空格或乱码
df.columns = df.columns.str.strip()

# 按作者等级和分类统计作品数量
heatmap_data = (
    df.groupby(["作者等级", "分类"])
    .size()
    .reset_index(name="作品数量")
)

# 提取热力图所需的数据
authors_levels = sorted(heatmap_data["作者等级"].unique().tolist())
categories = sorted(heatmap_data["分类"].unique().tolist())
data_matrix = [
    [authors_levels.index(row["作者等级"]), categories.index(row["分类"]), row["作品数量"]]
    for _, row in heatmap_data.iterrows()
]

# 创建热力图
heatmap = (
    HeatMap()
    .add_xaxis(categories)  # 分类作为横轴
    .add_yaxis("作者等级", authors_levels, data_matrix)  # 作者等级作为纵轴
    .set_global_opts(
        title_opts=opts.TitleOpts(title="作者等级与分类的作品数量热力图"),
        tooltip_opts=opts.TooltipOpts(
            is_show=True, formatter="分类: {b0}<br>作者等级: {a0}<br>作品数量: {c}"
        ),
        visualmap_opts=opts.VisualMapOpts(
            max_=heatmap_data["作品数量"].max(),
            min_=0,
            orient="vertical",
            pos_left="90%",
            pos_top="10%",
        ),
        xaxis_opts=opts.AxisOpts(name="分类", axislabel_opts={"rotate": 30}),
        yaxis_opts=opts.AxisOpts(name="作者等级"),
    )
)

# 保存图表
output_file = "author_category_heatmap.html"
heatmap.render(output_file)
print(f"图表已保存为 {output_file}，打开即可查看。")
